{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import json\n",
    "import random\n",
    "import re\n",
    "from src.functions import *\n",
    "import random\n",
    "import datetime\n",
    "\n",
    "import absl.logging #prevent checkpoint warnings while training\n",
    "absl.logging.set_verbosity(absl.logging.ERROR)\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obj_library = {}\n",
    "with open(\"imagenet1000_clsidx_to_labels.txt\") as f:\n",
    "    obj_library = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_obj_in_text(obj_names,text):\n",
    "\n",
    "    obj = \"\"\n",
    "    for w in re.split(\"\\s|(?<!\\d)[,.](?!\\d)\",text):\n",
    "        if w in obj_names:\n",
    "            obj = w\n",
    "    return obj\n",
    "def change_obj_names(obj_names,text):\n",
    "    new_objs_class = random.sample(obj_library.keys(),len(obj_names))\n",
    "    new_objs = [random.choice(obj_library[o]) for o in new_objs_class]\n",
    "\n",
    "    obj = find_obj_in_text(obj_names,text)\n",
    "    if obj != \"\":\n",
    "        new_obj = new_objs[obj_names.index(obj)]\n",
    "        text = text.replace(obj,new_obj)\n",
    "        obj = new_obj\n",
    "    return new_objs, obj, text\n",
    "\n",
    "def load_data(data_dir = \"../data/train2\", change_names =False):\n",
    "    \n",
    "    json_files = [pos_json for pos_json in os.listdir(data_dir) if pos_json.endswith('.json')]\n",
    "    data = []\n",
    "    for index, js in enumerate(json_files):\n",
    "        with open(os.path.join(data_dir, js)) as json_file:\n",
    "            json_text = json.load(json_file)\n",
    "\n",
    "            #TODO: change text extraction\n",
    "            text = js.split(\"_\", 1)[1][:-8] \n",
    "\n",
    "            #TODO: change in the normalization\n",
    "            objs_x = np.array(json_text[\"meta\"][\"o_center_x\"])/json_text[\"meta\"][\"width\"]\n",
    "            objs_y = np.array(json_text[\"meta\"][\"o_center_y\"])/json_text[\"meta\"][\"height\"]\n",
    "\n",
    "\n",
    "            #replace objs for random ones, avoids repetive obj names\n",
    "            obj_names = json_text[\"meta\"][\"obj_names\"]\n",
    "            obj = \"\"\n",
    "\n",
    "            if change_names:\n",
    "                new_objs, obj, text = change_obj_names(obj_names,text)\n",
    "            else:\n",
    "                new_objs = obj_names\n",
    "                obj = find_obj_in_text(obj_names,text)\n",
    "\n",
    "            \n",
    "            #TODO: add radius\n",
    "            data.append({\"input_traj\":json_text[\"input_traj\"],\n",
    "                        \"output_traj\":json_text[\"output_traj\"],\n",
    "                        \"text\":text,\n",
    "                        \"obj_names\":new_objs,\n",
    "                        \"obj_poses\":np.stack([objs_x,objs_y],axis = 0),\n",
    "                        \"obj_in_text\":obj\n",
    "                        })\n",
    "            # print(data)\n",
    "            # break\n",
    "\n",
    "    return data\n",
    "data = load_data(data_dir = \"../data/train\", change_names=True) + load_data(data_dir = \"../data/train2/train2\", change_names=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src.motion_refiner import Motion_refiner\n",
    "\n",
    "traj_n = 10\n",
    "mr = Motion_refiner(traj_n = traj_n)\n",
    "\n",
    "## ------- processed data -------\n",
    "X,Y = mr.prepare_data(data,deltas=True)\n",
    "print(\"X: \",X.shape)\n",
    "print(\"Y: \",Y.shape)\n",
    "\n",
    "## ------- save pre processed data -------\n",
    "mr.save_XY(X, Y, x_name=\"X_delta_new_names\",y_name=\"Y_delta_new_names\")\n",
    "mr.save_data(data,data_name=\"data_delta_new_names\")\n",
    "\n",
    "# ------- load data --------\n",
    "# X_, Y_ = mr.load_XY(x_name=\"X_delta_new_names\",y_name=\"Y_delta_new_names\")\n",
    "# data_ = mr.load_data(data_name=\"data_delta_new_names\")\n",
    "# feature_indices, obj_sim_indices, obj_poses_indices, traj_indices = mr.get_indices()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
